package profiler

import (
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"time"

	"github.com/gin-gonic/gin"
)

// Profiler 性能分析器
type Profiler struct {
	cpuFile   *os.File
	memFile   *os.File
	goroutine bool
	block     bool
}

// NewProfiler 创建性能分析器
func NewProfiler(options ...Option) (*Profiler, error) {
	p := &Profiler{}
	for _, opt := range options {
		opt(p)
	}
	return p, nil
}

// Option 配置选项
type Option func(*Profiler)

// WithCPUProfile 启用CPU性能分析
func WithCPUProfile(filename string) Option {
	return func(p *Profiler) {
		f, err := os.Create(filename)
		if err != nil {
			panic(err)
		}
		p.cpuFile = f
	}
}

// WithMemProfile 启用内存性能分析
func WithMemProfile(filename string) Option {
	return func(p *Profiler) {
		f, err := os.Create(filename)
		if err != nil {
			panic(err)
		}
		p.memFile = f
	}
}

// WithGoroutineProfile 启用goroutine分析
func WithGoroutineProfile() Option {
	return func(p *Profiler) {
		p.goroutine = true
	}
}

// WithBlockProfile 启用阻塞分析
func WithBlockProfile() Option {
	return func(p *Profiler) {
		p.block = true
	}
}

// Start 开始性能分析
func (p *Profiler) Start() {
	if p.cpuFile != nil {
		pprof.StartCPUProfile(p.cpuFile)
	}
	if p.block {
		runtime.SetBlockProfileRate(1)
	}
}

// Stop 停止性能分析
func (p *Profiler) Stop() {
	if p.cpuFile != nil {
		pprof.StopCPUProfile()
		p.cpuFile.Close()
	}
	if p.memFile != nil {
		runtime.GC()
		pprof.WriteHeapProfile(p.memFile)
		p.memFile.Close()
	}
}

// RegisterHTTPHandlers 注册HTTP处理器
func (p *Profiler) RegisterHTTPHandlers(r *gin.Engine) {
	r.GET("/debug/pprof/profile", p.cpuProfile)
	r.GET("/debug/pprof/heap", p.heapProfile)
	r.GET("/debug/pprof/goroutine", p.goroutineProfile)
	r.GET("/debug/pprof/block", p.blockProfile)
}

func (p *Profiler) cpuProfile(c *gin.Context) {
	duration := 30 * time.Second
	if s := c.Query("seconds"); s != "" {
		fmt.Sscanf(s, "%d", &duration)
	}

	f, err := os.Create("cpu.pprof")
	if err != nil {
		c.String(500, err.Error())
		return
	}
	pprof.StartCPUProfile(f)
	time.Sleep(duration)
	pprof.StopCPUProfile()
	f.Close()
	c.File("cpu.pprof")
}

func (p *Profiler) heapProfile(c *gin.Context) {
	f, err := os.Create("heap.pprof")
	if err != nil {
		c.String(500, err.Error())
		return
	}
	pprof.WriteHeapProfile(f)
	f.Close()
	c.File("heap.pprof")
}

func (p *Profiler) goroutineProfile(c *gin.Context) {
	f, err := os.Create("goroutine.pprof")
	if err != nil {
		c.String(500, err.Error())
		return
	}
	pprof.Lookup("goroutine").WriteTo(f, 1)
	f.Close()
	c.File("goroutine.pprof")
}

func (p *Profiler) blockProfile(c *gin.Context) {
	f, err := os.Create("block.pprof")
	if err != nil {
		c.String(500, err.Error())
		return
	}
	pprof.Lookup("block").WriteTo(f, 1)
	f.Close()
	c.File("block.pprof")
} 